/**
 * @author Yunfeng Bai
 *
 * 2013-6-25
 */
package simulation;

import cern.colt.matrix.DoubleMatrix1D;
import cern.colt.matrix.DoubleMatrix2D;
import cern.colt.matrix.impl.DenseDoubleMatrix1D;
import cern.colt.matrix.impl.SparseDoubleMatrix2D;
import cern.colt.matrix.linalg.Algebra;
import cern.colt.matrix.linalg.EigenvalueDecomposition;
import cern.colt.matrix.linalg.Property;
import cern.jet.math.Functions;

/***
 * Solve the equation Ax + b = w, w >= 0, x >= 0, x^T * w = 0, with restriction lo <= x <= hi, using projected Gauss-Seidel method.
 * See "Velocity-Based Shock Propagation for Multibody Dynamics Animation", Erleben, K., equation 28.
 */
public class ProjectedGaussSeidel
{
	protected DoubleMatrix1D x, newX;
	protected int size;
	
	public class PGSResult
	{
		public DoubleMatrix1D result;
		public int iterationProcessed;
	};
	
	public PGSResult solve(DoubleMatrix2D a, DoubleMatrix1D b,
			DoubleMatrix1D lo, DoubleMatrix1D hi, int maxIteration, double threshold)
	{
		if (a.columns() != a.rows())
			throw new RuntimeException("The matrix A must be squared.");
		
		size = a.columns();
		x = new DenseDoubleMatrix1D(size);
		newX = new DenseDoubleMatrix1D(size);
		DoubleMatrix1D newX1 = new DenseDoubleMatrix1D(size);
		
		// Build B and C.
		DoubleMatrix2D B = new SparseDoubleMatrix2D(size, size);
		DoubleMatrix2D C = new SparseDoubleMatrix2D(size, size);
		for (int i = 0; i < size; i++)
			for (int j = 0; j <= i; j++)
			{
				double v = a.get(i, j);
				if (v != 0)
					B.set(i, j, v);
			}
		for (int i = 0; i < size; i++)
			for (int j = i + 1; j < size; j++)
			{
				double v = a.get(i, j);
				if (v != 0)
					C.set(i, j, v);
			}
		
		// Check if (B, C) is a regular splitting.
		DoubleMatrix2D BmC = new SparseDoubleMatrix2D(size, size);
		BmC.assign(B);
		BmC.assign(C, Functions.minus);
		EigenvalueDecomposition ed = new EigenvalueDecomposition(BmC);
		for (int i = 0; i < size; i++)
			if (ed.getRealEigenvalues().get(i) < 0)
				throw new RuntimeException("B - C is not positive semi-definite");
		
		boolean isSingular = (new Property(1e-12)).isSingular(B);
//		if (isSingular)
//			throw new RuntimeException("B is not invertible!");
		B = (new Algebra()).inverse(B);
		
		int it;
		DoubleMatrix1D u = new DenseDoubleMatrix1D(size);
		double lastF = Double.MAX_VALUE;
		double f = 0;
		for (it = 0; it < maxIteration; it++)
		{
			// Matrix form.
			C.zMult(x, u);
			u.assign(b, Functions.plus);
			u.assign(Functions.mult(-1));
			B.zMult(u, newX);
			
			// Item form.
			for (int i = 0; i < size; i++)
			{
				double v = -b.get(i);
				for (int j = 0; j < i; j++)
					v -= a.get(i, j) * newX1.get(j);
				for (int j = i + 1; j < size; j++)
					v -= a.get(i, j) * x.get(j);
				newX1.set(i, v / a.get(i, i));
			}
			
			for (int i = 0; i < size; i++)
				if (Math.abs(newX1.get(i) - newX.get(i)) > 1e-10)
					throw new RuntimeException("x_{k+1} is different!");
			
			DoubleMatrix1D temp = x;
			x = newX;
			newX = temp;
			
			f = CalculateF(a, b, x);
			if (lastF - f < 0)
			{
//				throw new RuntimeException("F is increasing");
			}
			lastF = f;
			
			// Projection.
			for (int i = 0; i < size; i++)
				x.set(i, Math.min(hi.get(i), Math.max(lo.get(i), x.get(i))));
			
			// Check stopping criteria (TODO we need a better criteria).
			double maxDiff = Double.NEGATIVE_INFINITY;
			for (int i = 0; i < size; i++)
			{
				double diff = Math.abs(x.get(i) - newX.get(i));
				if (diff > maxDiff)
					maxDiff = diff;
			}
			if (maxDiff < threshold)
				break;
			
			// Debug: calculate the quadratic problem f(x).
		}
		PGSResult result = new PGSResult();
		result.result = x;
		result.iterationProcessed = it;
		return result;
	}
	
	private double CalculateF(DoubleMatrix2D A, DoubleMatrix1D b, DoubleMatrix1D x)
	{
		double v = 0;
		v += x.zDotProduct(b);
		DoubleMatrix1D Ax = new DenseDoubleMatrix1D(A.rows());
		A.zMult(x, Ax);
		v += x.zDotProduct(Ax) / 2;
		return v;
	}
	
	/**
	 * Solve with a start value. A good start value may improve performance.
	 */
	public PGSResult solve(DoubleMatrix2D a, DoubleMatrix1D b, DoubleMatrix1D lo, DoubleMatrix1D hi, int maxIteration, double threshold,
			DoubleMatrix1D start)
	{
		x = start;
		return solve(a, b, lo, hi, maxIteration, threshold);
	}
}
